import os
from abc import abstractmethod
from services.ProcessDoc import PDFLoader, TextLoader, Doc, TextToList
from services.Predictor import Predictor


class Task(object):
    def __init__(self, priority: int = 0, repo_file=None, callBack=None):
        self.priority = priority
        self.callBack = callBack
        self.repo_file = repo_file

    @abstractmethod
    def execute(self):
        pass

    def setCallBack(self, callback):
        self.callBack = callback

    @abstractmethod
    def get_result(self):
        pass

    def __lt__(self, other):
        """ 定义<比较操作符。 """
        return self.priority < other.priority


class FilePredictTask(Task):
    def __init__(self, repo_file, priority: int = 0, callBack=None):
        super(FilePredictTask, self).__init__(priority=priority, repo_file=repo_file, callBack=callBack)
        file_loader = {
            'pdf': PDFLoader,
            'txt': TextLoader
        }
        self.file_path = repo_file.fullname
        if os.path.exists(self.file_path):
            file_type = self.file_path.split(".")[-1]
            self.Loader = file_loader[file_type]
        else:
            raise "所给上传的文件路径不对 filepredicttask"
        self.doc = None
        self.predictor = Predictor()

    def execute(self):
        print("------execute file -----")
        self.doc = self.Loader(self.file_path).load()
        content, _ = self.doc.get_content()
        text_list = TextToList(content).getList()
        self.predictor \
            .putinData(sentence_list=text_list, use_dataset_file=False) \
            .predict()
        result = self.predictor.get_result()
        self.doc.set_triple_dict(result)  # {"predict_triple": ..., "text":...}
        self.callBack(result, self.repo_file)

    def get_result(self):
        return self.doc


class DatasetPredictTask(Task):
    def __init__(self, repo_file, priority: int = 0, callBack=None):
        super(DatasetPredictTask, self).__init__(priority=priority, repo_file=repo_file, callBack=callBack)
        self.file_path = repo_file.fullname
        if not os.path.exists(self.file_path):
            raise "所给上传的文件路径不对 DatasetPredictTask"
        self.doc = None
        self.predictor = Predictor()

    def execute(self):
        print("------execute datasets -----")
        self.doc = Doc(self.file_path)
        self.predictor \
            .putinData(use_dataset_file=True, dataset_file=self.file_path) \
            .predict()
        result = self.predictor.get_result()
        self.doc.set_triple_dict(result)  # {"predict_triple": ..., "text":...}
        self.callBack(result, self.repo_file)

    def get_result(self):
        return self.doc


if __name__ == '__main__':
    # task = FilePredictTask(
    #     file_path="tes.txt")
    task = DatasetPredictTask(
        "D:\Project\Django_KG\KnowledgeGraph\services\Predictor\Datades\seq2umt_ops\\new_test_data.json")
    task.execute()
    result = task.get_result()
    print(result)
